#!/usr/bin/env python3
"""
MCP协议服务器 - 完全符合MCP协议规范
"""

import json
import sys
import os
import locale

# 在文件开头添加
os.environ['PYTHONIOENCODING'] = 'utf-8'

# 修改stdin/stdout编码设置
if sys.platform.startswith('win'):
    import codecs
    try:
        sys.stdout.reconfigure(encoding='utf-8', errors='replace')
        sys.stderr.reconfigure(encoding='utf-8', errors='replace')
        sys.stdin.reconfigure(encoding='utf-8', errors='replace')
    except:
        sys.stdout = codecs.getwriter('utf-8')(sys.stdout.detach())
        sys.stderr = codecs.getwriter('utf-8')(sys.stderr.detach())
        sys.stdin = codecs.getreader('utf-8')(sys.stdin.detach())

# 修复导入路径问题
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.insert(0, parent_dir)

try:
    from mcp.start_mcp import SimpleSSHServer
except ImportError:
    sys.path.insert(0, current_dir)
    from start_mcp import SimpleSSHServer

class MCPProtocolServer:
    """MCP协议服务器"""
    
    def __init__(self):
        self.ssh_server = SimpleSSHServer()
        self.initialized = False
        
        # MCP工具定义
        self.tools = {
            "list_servers": {
                "description": "列出所有配置的服务器",
                "inputSchema": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            },
            "execute_command": {
                "description": "在指定服务器上执行SSH命令",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "server_name": {"type": "string", "description": "服务器名称"},
                        "command": {"type": "string", "description": "要执行的命令"}
                    },
                    "required": ["server_name", "command"]
                }
            },
            "add_server": {
                "description": "添加新的服务器配置",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "name": {"type": "string", "description": "服务器名称"},
                        "ip": {"type": "string", "description": "IP地址"},
                        "port": {"type": "integer", "description": "SSH端口", "default": 22},
                        "username": {"type": "string", "description": "用户名"},
                        "auth_type": {"type": "string", "enum": ["password", "key"], "description": "认证方式"},
                        "password": {"type": "string", "description": "密码（当auth_type为password时）"},
                        "key_path": {"type": "string", "description": "密钥路径（当auth_type为key时）"},
                        "description": {"type": "string", "description": "描述"}
                    },
                    "required": ["name", "ip", "username", "auth_type"]
                }
            },
            "delete_server": {
                "description": "删除服务器配置",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "server_name": {"type": "string", "description": "服务器名称"}
                    },
                    "required": ["server_name"]
                }
            },
            "get_server_info": {
                "description": "获取指定服务器的详细信息",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "server_name": {"type": "string", "description": "服务器名称"}
                    },
                    "required": ["server_name"]
                }
            },
            "test_connection": {
                "description": "测试与指定服务器的SSH连接",
                "inputSchema": {
                    "type": "object",
                    "properties": {
                        "server_name": {"type": "string", "description": "服务器名称"}
                    },
                    "required": ["server_name"]
                }
            }
        }
    
    def safe_encode_text(self, text):
        """安全编码文本，确保中文正确显示"""
        if isinstance(text, str):
            # 确保文本是UTF-8编码
            try:
                # 如果文本包含非ASCII字符，确保正确编码
                text.encode('utf-8').decode('utf-8')
                return text
            except UnicodeError:
                # 如果编码有问题，尝试修复
                return text.encode('utf-8', errors='replace').decode('utf-8')
        return str(text)
    
    def send_response(self, response):
        """发送MCP响应"""
        try:
            # 使用ensure_ascii=False确保中文正确显示
            json_response = json.dumps(response, ensure_ascii=False, separators=(',', ':'))
            print(json_response, flush=True)
        except Exception as e:
            # 如果JSON序列化失败，尝试使用ASCII编码
            try:
                json_response = json.dumps(response, ensure_ascii=True, separators=(',', ':'))
                print(json_response, flush=True)
            except Exception as e2:
                # 最后的备用方案
                error_response = {
                    "jsonrpc": "2.0",
                    "id": response.get("id"),
                    "error": {
                        "code": -32603,
                        "message": f"响应序列化失败: {str(e2)}"
                    }
                }
                print(json.dumps(error_response, ensure_ascii=True), flush=True)
    
    def send_notification(self, method, params=None):
        """发送MCP通知"""
        notification = {
            "jsonrpc": "2.0",
            "method": method
        }
        if params:
            notification["params"] = params
        self.send_response(notification)
    
    def handle_initialize(self, request_id, params):
        """处理初始化请求"""
        response = {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {}},
                "serverInfo": {
                    "name": "mcp-ssh-tools",
                    "version": "1.0.0"
                }
            }
        }
        self.send_response(response)
        self.send_notification("notifications/ready")
        self.initialized = True
    
    def handle_tools_list(self, request_id):
        """处理工具列表请求"""
        response = {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "tools": [
                    {
                        "name": tool_name,
                        "description": tool_info["description"],
                        "inputSchema": tool_info["inputSchema"]
                    }
                    for tool_name, tool_info in self.tools.items()
                ]
            }
        }
        self.send_response(response)
    
    def handle_tools_call(self, request_id, params):
        """处理工具调用请求"""
        tool_name = params.get("name")
        arguments = params.get("arguments", {})
        
        try:
            if tool_name == "list_servers":
                result = self.ssh_server.list_servers()
                if result.get("success"):
                    servers = result.get('servers', [])
                    if servers:
                        server_lines = []
                        for s in servers:
                            name = self.safe_encode_text(s.get('name', ''))
                            ip = self.safe_encode_text(s.get('ip', ''))
                            port = s.get('port', 22)
                            server_lines.append(f"- {name} ({ip}:{port})")
                        
                        content = f"找到 {len(servers)} 台服务器:\n" + "\n".join(server_lines)
                    else:
                        content = "当前没有配置任何服务器"
                else:
                    error_msg = self.safe_encode_text(result.get('error', '未知错误'))
                    content = f"获取服务器列表失败: {error_msg}"
                
                response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {
                        "content": [{"type": "text", "text": content}]
                    }
                }
            
            elif tool_name == "execute_command":
                server_name = arguments.get("server_name")
                command = arguments.get("command")
                result = self.ssh_server.execute_command(server_name, command)
                
                if result.get("success"):
                    output = self.safe_encode_text(result.get('output', '无输出'))
                    content = f"命令执行成功:\n{output}"
                else:
                    error_msg = self.safe_encode_text(result.get('error', '未知错误'))
                    content = f"命令执行失败: {error_msg}"
                
                response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {
                        "content": [{"type": "text", "text": content}]
                    }
                }
            
            elif tool_name == "add_server":
                result = self.ssh_server.add_server(arguments)
                if result.get("success"):
                    message = self.safe_encode_text(result.get("message", "服务器添加成功"))
                    content = message
                else:
                    error_msg = self.safe_encode_text(result.get('error', '未知错误'))
                    content = f"添加服务器失败: {error_msg}"
                
                response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {
                        "content": [{"type": "text", "text": content}]
                    }
                }
            
            elif tool_name == "delete_server":
                server_name = arguments.get("server_name")
                result = self.ssh_server.delete_server(server_name)
                if result.get("success"):
                    message = self.safe_encode_text(result.get("message", "服务器删除成功"))
                    content = message
                else:
                    error_msg = self.safe_encode_text(result.get('error', '未知错误'))
                    content = f"删除服务器失败: {error_msg}"
                
                response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {
                        "content": [{"type": "text", "text": content}]
                    }
                }
            
            elif tool_name == "get_server_info":
                server_name = arguments.get("server_name")
                result = self.ssh_server.get_server_info(server_name)
                if result.get("success"):
                    server = result.get("server", {})
                    name = self.safe_encode_text(server.get('name', ''))
                    ip = self.safe_encode_text(server.get('ip', ''))
                    port = server.get('port', 22)
                    username = self.safe_encode_text(server.get('username', ''))
                    auth_type = self.safe_encode_text(server.get('auth_type', ''))
                    
                    content = f"服务器信息:\n名称: {name}\nIP: {ip}\n端口: {port}\n用户名: {username}\n认证方式: {auth_type}"
                else:
                    error_msg = self.safe_encode_text(result.get('error', '未知错误'))
                    content = f"获取服务器信息失败: {error_msg}"
                
                response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {
                        "content": [{"type": "text", "text": content}]
                    }
                }
            
            elif tool_name == "test_connection":
                server_name = arguments.get("server_name")
                result = self.ssh_server.test_connection(server_name)
                if result.get("success"):
                    message = self.safe_encode_text(result.get("message", "连接测试成功"))
                    content = message
                else:
                    error_msg = self.safe_encode_text(result.get('error', '未知错误'))
                    content = f"连接测试失败: {error_msg}"
                
                response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {
                        "content": [{"type": "text", "text": content}]
                    }
                }
            
            else:
                response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "error": {
                        "code": -32602,
                        "message": f"未知工具: {tool_name}"
                    }
                }
            
            self.send_response(response)
            
        except Exception as e:
            error_response = {
                "jsonrpc": "2.0",
                "id": request_id,
                "error": {
                    "code": -32603,
                    "message": f"内部错误: {str(e)}"
                }
            }
            self.send_response(error_response)
    
    def run(self):
        """运行MCP协议服务器"""
        try:
            while True:
                line = input()
                if not line.strip():
                    continue
                
                try:
                    request = json.loads(line)
                except json.JSONDecodeError:
                    continue
                
                method = request.get("method")
                request_id = request.get("id")
                params = request.get("params", {})
                
                if method == "initialize":
                    self.handle_initialize(request_id, params)
                elif method == "tools/list":
                    self.handle_tools_list(request_id)
                elif method == "tools/call":
                    self.handle_tools_call(request_id, params)
                elif method == "ping":
                    response = {
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "result": {}
                    }
                    self.send_response(response)
                else:
                    response = {
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "error": {
                            "code": -32601,
                            "message": f"未知方法: {method}"
                        }
                    }
                    self.send_response(response)
                    
        except (EOFError, KeyboardInterrupt):
            pass

if __name__ == "__main__":
    server = MCPProtocolServer()
    server.run() 